#ifndef ASTNODE_H
#define ASTNODE_H

#include <string>
#include <vector>
#include <memory>
#include <iostream>
#include <map>

// 前向声明
class Module;
class Function;
class BasicBlock;
class IRBuilder;
class Value;
class Type;

// AST节点类型枚举
enum class NodeType {
    // 程序结构
    PROGRAM,
    COMP_UNIT,
    
    // 声明
    CONST_DECL,
    VAR_DECL,
    CONST_DEF,
    VAR_DEF,
    
    // 函数
    FUNC_DEF,
    FUNC_FPARAM,
    FUNC_FPARAMS,
    FUNC_RPARAM,
    FUNC_RPARAMS,
    
    // 语句
    BLOCK,
    ASSIGN_STMT,
    EXPR_STMT,
    RETURN_STMT,
    IF_STMT,
    
    // 表达式
    BINARY_EXPR,
    UNARY_EXPR,
    CALL_EXPR,
    IDENT_EXPR,
    LITERAL_EXPR,
    PAREN_EXPR,
    
    // 基本类型
    BASIC_TYPE,
    
    // 其他
    INIT_VAL,
    CONST_INIT_VAL
};

// AST节点基类
class ASTNode {
public:
    std::string node_type;
    std::string value;
    std::vector<std::unique_ptr<ASTNode>> children;
    
    ASTNode(const std::string& type) : node_type(type) {}
    ASTNode(const std::string& type, const std::string& val) : node_type(type), value(val) {}
    virtual ~ASTNode() = default;
    
    void add_child(std::unique_ptr<ASTNode> child) {
        children.push_back(std::move(child));
    }
    
    // 调试打印函数
    void print_ast(int indent = 0) const {
        for (int i = 0; i < indent; ++i) std::cout << "  ";
        std::cout << node_type;
        if (!value.empty()) {
            std::cout << " (" << value << ")";
        }
        std::cout << std::endl;
        
        for (const auto& child : children) {
            child->print_ast(indent + 1);
        }
    }
    
    // 生成IR的虚函数
    virtual Value* generate_ir(IRBuilder* builder) = 0;
};

// 程序节点
class ProgramNode : public ASTNode {
public:
    ProgramNode() : ASTNode("Program") {}
    Value* generate_ir(IRBuilder* builder) override;
};

// 编译单元节点
class CompUnitNode : public ASTNode {
public:
    CompUnitNode() : ASTNode("CompUnit") {}
    Value* generate_ir(IRBuilder* builder) override;
};

// 函数定义节点
class FuncDefNode : public ASTNode {
public:
    FuncDefNode(const std::string& name) : ASTNode("FuncDef", name) {}
    Value* generate_ir(IRBuilder* builder) override;
};

// 函数形参节点
class FuncFParamNode : public ASTNode {
public:
    FuncFParamNode(const std::string& name) : ASTNode("FuncFParam", name) {}
    Value* generate_ir(IRBuilder* builder) override;
};

// 函数形参列表节点
class FuncFParamsNode : public ASTNode {
public:
    FuncFParamsNode() : ASTNode("FuncFParams") {}
    Value* generate_ir(IRBuilder* builder) override;
};

// 基本类型节点
class BasicTypeNode : public ASTNode {
public:
    BasicTypeNode(const std::string& type) : ASTNode("BasicType", type) {}
    Value* generate_ir(IRBuilder* builder) override;
    Type* get_type(Module* module);
};

// 常量声明节点
class ConstDeclNode : public ASTNode {
public:
    ConstDeclNode() : ASTNode("ConstDecl") {}
    Value* generate_ir(IRBuilder* builder) override;
};

// 变量声明节点
class VarDeclNode : public ASTNode {
public:
    VarDeclNode() : ASTNode("VarDecl") {}
    Value* generate_ir(IRBuilder* builder) override;
};

// 常量定义节点
class ConstDefNode : public ASTNode {
public:
    ConstDefNode(const std::string& name) : ASTNode("ConstDef", name) {}
    Value* generate_ir(IRBuilder* builder) override;
};

// 变量定义节点
class VarDefNode : public ASTNode {
public:
    VarDefNode(const std::string& name) : ASTNode("VarDef", name) {}
    Value* generate_ir(IRBuilder* builder) override;
};

// 初始化值节点
class InitValNode : public ASTNode {
public:
    InitValNode() : ASTNode("InitVal") {}
    Value* generate_ir(IRBuilder* builder) override;
};

// 常量初始化值节点
class ConstInitValNode : public ASTNode {
public:
    ConstInitValNode() : ASTNode("ConstInitVal") {}
    Value* generate_ir(IRBuilder* builder) override;
};

// 代码块节点
class BlockNode : public ASTNode {
public:
    BlockNode() : ASTNode("Block") {}
    Value* generate_ir(IRBuilder* builder) override;
};

// 赋值语句节点
class AssignStmtNode : public ASTNode {
public:
    AssignStmtNode() : ASTNode("AssignStmt") {}
    Value* generate_ir(IRBuilder* builder) override;
};

// 表达式语句节点
class ExprStmtNode : public ASTNode {
public:
    ExprStmtNode() : ASTNode("ExprStmt") {}
    Value* generate_ir(IRBuilder* builder) override;
};

// 返回语句节点
class ReturnStmtNode : public ASTNode {
public:
    ReturnStmtNode() : ASTNode("ReturnStmt") {}
    Value* generate_ir(IRBuilder* builder) override;
};

// if语句节点
class IfStmtNode : public ASTNode {
public:
    IfStmtNode() : ASTNode("IfStmt") {}
    Value* generate_ir(IRBuilder* builder) override;
};

// 二元表达式节点
class BinaryExprNode : public ASTNode {
public:
    std::string op;
    BinaryExprNode(const std::string& operator_) : ASTNode("BinaryExpr"), op(operator_) {}
    Value* generate_ir(IRBuilder* builder) override;
};

// 一元表达式节点
class UnaryExprNode : public ASTNode {
public:
    std::string op;
    UnaryExprNode(const std::string& operator_) : ASTNode("UnaryExpr"), op(operator_) {}
    Value* generate_ir(IRBuilder* builder) override;
};

// 函数调用表达式节点
class CallExprNode : public ASTNode {
public:
    CallExprNode(const std::string& func_name) : ASTNode("CallExpr", func_name) {}
    Value* generate_ir(IRBuilder* builder) override;
};

// 标识符表达式节点
class IdentExprNode : public ASTNode {
public:
    IdentExprNode(const std::string& name) : ASTNode("IdentExpr", name) {}
    Value* generate_ir(IRBuilder* builder) override;
};

// 字面量表达式节点
class LiteralExprNode : public ASTNode {
public:
    std::string literal_type;
    LiteralExprNode(const std::string& type, const std::string& val) : ASTNode("LiteralExpr", val) {
        // 解析类型信息，例如 "int: 0" -> type="int", value="0"
        if (type.find(':') != std::string::npos) {
            size_t pos = type.find(':');
            this->literal_type = type.substr(0, pos);
            this->value = type.substr(pos + 2); // 跳过": "
        } else {
            this->literal_type = type;
        }
    }
    Value* generate_ir(IRBuilder* builder) override;
};

// 括号表达式节点
class ParenExprNode : public ASTNode {
public:
    ParenExprNode() : ASTNode("ParenExpr") {}
    Value* generate_ir(IRBuilder* builder) override;
};

#endif 